/************************************************************************
Copyright (c) 2010, 2013, C.R. All rights reserved. 

E-mail: crtoyou@gmail.com; crtoyou@163.com

You can redistribute it and/or modify it under the terms of the
GNU Lesser General Public License as published by the
Free Software Foundation version 3 of the License.

This source code is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program. if not, write to the Free Software Foundation, Inc., 
51 Franklin St, Fifth Floor, Boston, MA 02110-1301  USA
************************************************************************/

#ifndef __CR_HUFFMAN_CODEC_H__
#define __CR_HUFFMAN_CODEC_H__

#include <crfilter.h>
#include <crbitstrm.h>
#include <crhuffman.h>

namespace crhuffman {

template <size_t CODE_SPACE /*256|16*/>
class frequencies_wrap
{
public:
	typedef size_t frequencies[CODE_SPACE];

	frequencies freqs_;
	frequencies_wrap() {ZeroMemory(freqs_, sizeof(freqs_));}
	inline size_t& operator[](size_t i) {return freqs_[i];}
	inline const size_t& operator[](size_t i) const {return freqs_[i];}

	frequencies_wrap& operator += (const frequencies_wrap& other)
	{
		for (size_t i=0; i<CODE_SPACE; ++i)
			freqs_[i] += other[i];
		return *this;
	}

};

template <size_t CODE_SPACE /*256|16*/>
class code_table_wrap
{
public:
	typedef bitstring<CODE_SPACE> bitstring;
	typedef bitstring code_table[CODE_SPACE];

	code_table codes_;
	code_table_wrap() {ZeroMemory(codes_, sizeof(codes_));}
	inline bitstring& operator[](size_t i) {return codes_[i];}
	inline const bitstring& operator[](size_t i) const {return codes_[i];}

};

inline void freq_merge(frequencies_wrap<16>& wrap16, 
	const frequencies_wrap<256>& wrap256)
{
	size_t i;
	for (i=0; i<16; ++i) {
		wrap16[i] = 0;
	}
	for (i=0; i<256; ++i) {
		const byte il = i & 0xF;
		const byte ih = i >> 4;
		wrap16[il] += wrap256[i];
		wrap16[ih] += wrap256[i];
	}
}

inline void code_merge(code_table_wrap<256>& wrap256, 
	const code_table_wrap<16>& wrap16)
{
	for (size_t i=0; i<256; ++i) {
		const byte il = i & 0xF;
		const byte ih = i >> 4;
		if (wrap16[il].length() <= 0 || wrap16[ih].length() <= 0)
			continue;
		wrap256[i] = wrap16[il] + wrap16[ih];
	}
}

template <size_t CODE_SPACE /*256|16*/>
class huffman_header
{
public:
	typedef bitstring<CODE_SPACE> bitstring;
	typedef bitstring code_table[CODE_SPACE];
	typedef size_t frequencies[CODE_SPACE];

	// without flag and length
	// return whether or not be tiny_sizes
	bool calcsize(const code_table& codes, const frequencies& freqs, bit_size& bitlen)
	{
		bool tiny_sizes = true;
		for (size_t i=0; i<CODE_SPACE; ++i) {
			const byte size = (byte)codes[i].length();
			if (size & 0xF0)
				tiny_sizes = false;
			bitlen += size * (freqs[i] + 1); // within the palette
		}
		bitlen.oBy += tiny_sizes ? CODE_SPACE / 2 : CODE_SPACE;
		return tiny_sizes;
	}

	// without flag and length
	char writecodes(ostream_base& strm, const code_table& codes, bool tiny_sizes)
	{
		if (tiny_sizes) {
			for (int i=0; i<CODE_SPACE/2; ++i) {
				const int il = i * 2;
				const int ih = i * 2 + 1;
				const byte sizel = (byte)codes[il].length();
				const byte sizeh = (byte)codes[ih].length();
				strm.putbyte((sizel & 0x0F) | (sizeh << 4));
			}
			obstream bitstrm(strm);
			for (int i=0; i<CODE_SPACE/2; ++i) {
				const int il = i * 2;
				const int ih = i * 2 + 1;
				const byte sizel = (byte)codes[il].length();
				const byte sizeh = (byte)codes[ih].length();
				if (sizel > 0) {
					TRACE(_T("0x%02X(%d):"), il, sizel);
					DEBUGING(codes[il].dump());
					bitstrm.write<CODE_SPACE>(codes[il]);
					TRACE(_T("\n"));
				}
				if (sizeh > 0) {
					TRACE(_T("0x%02X(%d):"), ih, sizeh);
					DEBUGING(codes[ih].dump());
					bitstrm.write<CODE_SPACE>(codes[ih]);
					TRACE(_T("\n"));
				}
			}
			return bitstrm.final();
		}
		else {
			for (int i=0; i<CODE_SPACE; ++i) {
				const byte size = (byte)codes[i].length();
				strm.putbyte(size);
			}
			obstream bitstrm(strm);
			for (int i=0; i<CODE_SPACE; ++i) {
				const byte size = (byte)codes[i].length();
				if (size > 0) {
					TRACE(_T("0x%02X(%d):"), i, size);
					DEBUGING(codes[i].dump());
					bitstrm.write<CODE_SPACE>(codes[i]);
					TRACE(_T("\n"));
				}
			}
			return bitstrm.final();
		}
	}

	char readcodes(istream_base& strm, code_table& codes, bool tiny_sizes)
	{
		byte sizes[CODE_SPACE];
		if (tiny_sizes)
		{
			for (int i=0; i<CODE_SPACE/2; ++i)
				sizes[i] = strm.getbyte();
			ibstream bitstrm(strm);
			for (int i=0; i<CODE_SPACE/2; ++i)
			{
				const byte sizelh = sizes[i];
				const byte sizel = sizelh & 0x0F;
				const byte sizeh = sizelh >> 4;
				const int il = i * 2;
				const int ih = i * 2 + 1;
				if (sizel > 0) {
					TRACE(_T("0x%02X(%d):"), il, sizel);
					bitstrm.read<CODE_SPACE>(codes[il], sizel);
					DEBUGING(codes[il].dump());
					TRACE(_T("\n"));
				}
				if (sizeh > 0) {
					TRACE(_T("0x%02X(%d):"), ih, sizeh);
					bitstrm.read<CODE_SPACE>(codes[ih], sizeh);
					DEBUGING(codes[ih].dump());
					TRACE(_T("\n"));
				}
			}
			return bitstrm.final();
		}
		else 
		{
			for (int i=0; i<CODE_SPACE; ++i)
				sizes[i] = strm.getbyte();
			ibstream bitstrm(strm);
			for (int i=0; i<CODE_SPACE; ++i)
			{
				const byte size = sizes[i];
				if (size > 0) {
					TRACE(_T("0x%02X(%d):"), i, size);
					bitstrm.read<CODE_SPACE>(codes[i], size);
					DEBUGING(codes[i].dump());
					TRACE(_T("\n"));
				}
			}
			return bitstrm.final();
		}
	}

};

class huffman_builder 
	: public pusher_base<>
	, public frequencies_wrap<256>
{
	// tips:
	// when build tree, it add freq with 256 bit code space
	// freq of 16 bit code space can be merged from 256 bit code space
	// 16 bit code space must be spread to 256 bit code space before transform
	// for decode, it will rebuild 256 bit code space for 16 bit code space
	// the bytes of stream would not be divided any of time

public: 
	inline void push(byte ch)
	{
		++freqs_[ch];
	}

public:
	inline size_t total() const 
	{
		size_t size = 0;
		for (size_t i=0; i<256; ++i)
			size += freqs_[i];
		return size;
	}

	// some streams have been compressed already
	// thus the bytes are well-distributed
	// analyse makes the decision weather or not to encode
	// return false if no need to compress
	inline bool analyse(size_t size, size_t min_size = 64) const 
	{
		const double norm = 256.0 * (256.0 - 1.0) * 0.04 * 0.04;
		double bias = 0;
		if (size < min_size)
			return false;
		for (size_t i=0; i<256; ++i)
		{
			double diff = 1 - (double)freqs_[i] * 256 / size;
			bias += diff * diff;
		}
		return (bias > norm);
	}

};

class huffman_encoder 
	: public pusher_base<>
	, public code_table_wrap<256>
{
public:
	huffman_encoder() : m_strm(NULL) {}
	inline void setbstrm(obstream* strm) { m_strm = strm; }

	inline void push(byte ch)
	{
		ASSERT(m_strm);
		bitstring& bs = codes_[ch];
		if (bs.length() > 0) {
			m_strm->write(bs);
		}
	}

protected:
	obstream* m_strm;

};

template <class type_next_filter>
class huffman_decoder
	: public pusher_base<type_next_filter> // source filter may have no input
	, public huffman_tree<256>
{
public:
	huffman_decoder() : m_strm(NULL) {}
	inline void setbstrm(ibstream* strm) { m_strm = strm; }

	inline void bitscan(size_t bitlen)
	{ ASSERT(m_strm); m_strm->scan<huffman_decoder>(bitlen, this); }

	inline void bitscan(bit_addr endpos)
	{ ASSERT(m_strm); m_strm->scan<huffman_decoder>(endpos, this); }

protected:
	ibstream* m_strm;
	friend class ibstream;

private:
	inline void step(bool bi)
	{
		//ignore root node, but walk to next branch.  
		//pay attention to the 0 sized file
		//thus the root will have no child. 
		ASSERT(root_);
		ASSERT(root_->l || root_->r);
		ASSERT(cursor_ && !cursor_->leaf);
		cursor_ = (bi ? cursor_->r : cursor_->l);
		if (!cursor_)
			throw exception<huffman_tree>(step_damaged_stream);
		ASSERT(cursor_);
		if (cursor_->leaf) {
			ASSERT(!cursor_->l && !cursor_->r);
			pass(cursor_->key);
			cursor_ = root_;
		}
	}
};

}//namespace crhuffman

#endif//__CR_HUFFMAN_CODEC_H__